{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "532f5654",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6.318092346191406e-05\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "import torch\n",
    "\n",
    "x1 = torch.tensor([1 , 2], dtype = torch.int32, device=\"cuda\")\n",
    "\n",
    "x2 = torch.tensor([1 , 2], dtype = torch.int32, device=\"cuda\")\n",
    "\n",
    "c = x1 - x2\n",
    "\n",
    "c = x1 - x2\n",
    "\n",
    "c = x1 - x2\n",
    "\n",
    "\n",
    "t1 = time.time()\n",
    "\n",
    "c = x1 - x2\n",
    "\n",
    "t2 = time.time()\n",
    "\n",
    "print(t2 - t1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "77b1cc90",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- 初始数据 ---\n",
      "每个请求的序列长度 (seq_lens):\n",
      "tensor([3, 5])\n",
      "\n",
      "对应的物理Slots (padded):\n",
      "tensor([[10, 11, 12, -1, -1, -1],\n",
      "        [50, 51, 52, 53, 54, -1]])\n",
      "\n",
      "------------------------------\n",
      "\n",
      "--- 步骤 1: 创建掩码 ---\n",
      "批次中的最大序列长度 (max_len): 5\n",
      "\n",
      "用于比较的arange_tensor:\n",
      "tensor([0, 1, 2, 3, 4])\n",
      "\n",
      "扩展维度后的序列长度 (for broadcasting):\n",
      "tensor([[3],\n",
      "        [5]])\n",
      "\n",
      "生成的布尔掩码 (valid_token_mask):\n",
      "tensor([[ True,  True,  True, False, False],\n",
      "        [ True,  True,  True,  True,  True]])\n",
      "\n",
      "------------------------------\n",
      "\n",
      "--- 步骤 2: 应用掩码 ---\n",
      "与掩码形状对齐的Slots (裁剪后):\n",
      "tensor([[10, 11, 12, -1, -1],\n",
      "        [50, 51, 52, 53, 54]])\n",
      "\n",
      "被掩码提取并展平后的最终结果 (flattened_slots):\n",
      "tensor([10, 11, 12, 50, 51, 52, 53, 54])\n",
      "\n",
      "观察结果：\n",
      " - 请求0的有效slots [10, 11, 12] 被提取出来。\n",
      " - 请求1的有效slots [50, 51, 52, 53, 54] 被提取出来。\n",
      " - 它们被自动连接 (concatenate) 成了一个一维张量。\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "\n",
    "seq_lens = torch.tensor([3, 5])\n",
    "\n",
    "\n",
    "physical_slots_padded = torch.tensor([\n",
    "    [10, 11, 12, -1, -1, -1],  # 请求0的slots (长度3)\n",
    "    [50, 51, 52, 53, 54, -1]   # 请求1的slots (长度5)\n",
    "])\n",
    "\n",
    "print(\"--- 初始数据 ---\")\n",
    "print(f\"每个请求的序列长度 (seq_lens):\\n{seq_lens}\\n\")\n",
    "print(f\"对应的物理Slots (padded):\\n{physical_slots_padded}\\n\")\n",
    "print(\"-\" * 30 + \"\\n\")\n",
    "\n",
    "\n",
    "# --- 2. 创建掩码 (valid_token_mask) ---\n",
    "\n",
    "# a. 找到这批请求中的最大长度\n",
    "max_len = torch.max(seq_lens).item()\n",
    "print(\"--- 步骤 1: 创建掩码 ---\")\n",
    "print(f\"批次中的最大序列长度 (max_len): {max_len}\\n\") # 输出: 5\n",
    "\n",
    "# b. 创建一个从 0 到 max_len-1 的序列\n",
    "# [0, 1, 2, 3, 4]\n",
    "arange_tensor = torch.arange(max_len)\n",
    "print(f\"用于比较的arange_tensor:\\n{arange_tensor}\\n\")\n",
    "\n",
    "# c. 将 seq_lens 扩展一个维度，以便进行广播比较\n",
    "# [[3],\n",
    "#  [5]]\n",
    "seq_lens_unsqueezed = seq_lens.unsqueeze(1)\n",
    "print(f\"扩展维度后的序列长度 (for broadcasting):\\n{seq_lens_unsqueezed}\\n\")\n",
    "\n",
    "# d. 进行广播比较，生成掩码\n",
    "# 这个操作会比较 arange_tensor 的每一项和 seq_lens_unsqueezed 的每一行\n",
    "# [0,1,2,3,4] < 3  -> [T,T,T,F,F]\n",
    "# [0,1,2,3,4] < 5  -> [T,T,T,T,T]\n",
    "valid_token_mask = arange_tensor < seq_lens_unsqueezed\n",
    "print(f\"生成的布尔掩码 (valid_token_mask):\\n{valid_token_mask}\\n\")\n",
    "print(\"-\" * 30 + \"\\n\")\n",
    "\n",
    "\n",
    "# --- 3. 应用掩码并展平 ---\n",
    "\n",
    "# 在应用掩码之前，我们先确保被索引的张量和掩码的形状是对齐的。\n",
    "# 我们的 physical_slots_padded 有6列，而掩码只有5列，所以我们先裁剪一下。\n",
    "slots_to_be_masked = physical_slots_padded[:, :max_len]\n",
    "print(\"--- 步骤 2: 应用掩码 ---\")\n",
    "print(f\"与掩码形状对齐的Slots (裁剪后):\\n{slots_to_be_masked}\\n\")\n",
    "\n",
    "# 最关键的一步！\n",
    "# 使用布尔张量 valid_token_mask 来索引另一个张量。\n",
    "# PyTorch 会找到所有 `True` 位置对应的值，并把它们按行优先顺序放入一个新的1D张量中。\n",
    "flattened_slots = slots_to_be_masked[valid_token_mask]\n",
    "\n",
    "print(f\"被掩码提取并展平后的最终结果 (flattened_slots):\\n{flattened_slots}\\n\")\n",
    "print(\"观察结果：\")\n",
    "print(\" - 请求0的有效slots [10, 11, 12] 被提取出来。\")\n",
    "print(\" - 请求1的有效slots [50, 51, 52, 53, 54] 被提取出来。\")\n",
    "print(\" - 它们被自动连接 (concatenate) 成了一个一维张量。\")\n",
    "\n",
    "# 最终结果是 tensor([10, 11, 12, 50, 51, 52, 53, 54])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
